#!/usr/bin/env python
# encoding: utf-8
# Author: 潘世友（stpann@163.com）

import functools

import numpy as np
import pandas as pd


def check_if_in_range(para: str, x: str, inx: tuple | list) -> None:
    if isinstance(inx, list):
        if isinstance(x, int) or isinstance(x, float):
            if x not in inx:
                raise ValueError(f'variable {para}={x} is out of the options: {inx}')
        else:
            if x.lower() not in inx:
                raise ValueError(f'variable {para}={x} is out of the options: {inx}')
    elif isinstance(inx, tuple):
        if inx[0] is None:
            min_ = -np.inf
            max_ = inx[1]
        elif inx[1] is None:
            min_ = inx[0]
            max_ = np.inf
        else:
            min_, max_ = inx
        if x is not None and not (min_ <= x <= max_):
            raise ValueError(f'参数{para}={x}超出范围: {inx}')


def check_para(dict_xs: dict):
    """
    装饰器，不支持叠加多个装饰器，外层的装饰器在func.__code__.co_varnames读取的参数将是底层装饰器的参数
    """

    def wrapper(func):
        @functools.wraps(func)
        def __(*args, **kwargs):
            func_para = func.__code__.co_varnames[:func.__code__.co_argcount]
            for para, input_range in dict_xs.items():
                if para in func_para:
                    idx = func_para.index(para)
                    if para in kwargs.keys():
                        check_if_in_range(para, kwargs[para], input_range)
                    elif args.__len__() - 1 >= idx:
                        check_if_in_range(para, args[idx], input_range)
            return func(*args, **kwargs)

        return __

    return wrapper


def check_data(data: pd.DataFrame, list_columns: list[str], **kwargs):
    """ 检查数据是否满足要求 """
    dsn_col = data.columns

    out_col = []
    for col_ in list_columns:
        if col_ in kwargs:
            val_ = kwargs[col_]
            if val_ not in dsn_col:
                raise Exception('你输入的【%s】指标【%s】不存在于数据中' % (col_, val_))
        else:
            if col_ in dsn_col:
                val_ = col_
            else:
                raise Exception('请指定数据中【%s】指标' % col_)
        out_col.append(val_)

    data = data[out_col]
    data = data.rename(columns=dict(zip(out_col, list_columns)))
    return data


if __name__ == '__main__':
    # 以下仅对a,b作不能为None的规则限制
    @check_para({'a': {1, 2, 3}, 'd': [2, 3], 'e': [3, 6], 'c': (10, None)})
    def f(a, d=7, e=1, c=None):
        """ this is a test note """
        print(a, c, d, e)


    # f(1, d=None, c=11, e=3)
    help(f)

    f(a=2, d=2, e=5, c=11)  # 正常运行
